from langchain.prompts import PromptTemplate
from pydantic import BaseModel, Field
from datetime import datetime, timedelta
from LLM_Prompt.siliconflow import SiliconFlow
import json


# 使用SiliconFlow 或 OpenAI 接入大模型
llm = SiliconFlow()

# 关键信息提取 Prompt 模板
extraction_prompt = PromptTemplate(
    input_variables=["question", "current_date"],
    template=(
        "今天日期是 {current_date}。请从下面的问题中提取关键信息，并按照 JSON 的格式输出。\n"
        "问题: {question}\n"
        "所需的关键信息包括:\n"
        "1) 船名 (boat_name)：如果没有提及返回 \"未知\"\n"
        "2) 日期 (date)：\n"
        "   - 如果是单日时间（如“2025-03-17”）：\n"
        "       startTime：当日 00:00:00\n"
        "       endTime：次日 00:00:00\n"
        "   - 如果是日期范围（如“3月15到17日”）：\n"
        "       startTime：开始日期 00:00:00\n"
        "       endTime：结束日期的次日 00:00:00\n"
        "   - 未提及时间则返回 \"未知\"\n"
        "   - 相对时间（如“今天”、“昨天”或“前天”）需根据当前日期 {current_date} 计算\n"
        "输出格式示例：\n"
        "{{\"boat_name\": \"海洋一号\", \"startTime\": \"2025-03-17 00:00:00\", \"endTime\": \"2025-03-18 00:00:00\"}}\n"
        "严格按照输出格式示例输出结果，不需要其它内容"
    ),
)

# Pydantic 数据验证器
class ExtractedData(BaseModel):
    boat_name: str = Field(..., description="船的名字，未提及则为'未知'")
    startTime: str = Field(
        ...,
        description="起始时间，格式为 YYYY-MM-DD 00:00:00，未提及则为'未知'"
    )
    endTime: str = Field(
        ...,
        description="结束时间，格式为 YYYY-MM-DD 00:00:00，未提及则为'未知'"
    )

    @staticmethod
    def validate_output(response: str):
        try:
            # 处理可能存在的遗留格式
            if "date" in response:
                raise ValueError("检测到旧版date字段，请使用时间范围格式")

            data = ExtractedData.model_validate_json(response)
            return data.model_dump()
        except Exception as e:
            raise ValueError(f"数据解析错误: {e}")

# 清理模型输出工具
def clean_model_output(raw_output: str) -> str:
    """
    清理大模型的输出，去除额外的 Markdown 等格式符号，确保输出是合法 JSON。
    特别处理时间范围格式：
    1. 去除代码块标记
    2. 验证时间格式
    """
    # 去除 Markdown 代码块标记
    cleaned = raw_output.strip()
    if cleaned.startswith("```"):
        cleaned = cleaned[cleaned.find("{"): cleaned.rfind("}") + 1]

    # 验证时间格式
    try:
        data = json.loads(cleaned)
        # 检查必填字段
        if "startTime" not in data or "endTime" not in data:
            raise ValueError("缺少必要的时间字段")

        # 检查时间格式（允许"未知"或符合格式）
        for field in ["startTime", "endTime"]:
            value = data[field]
            if value != "未知":
                # 简单格式验证（可扩展为严格验证）
                if len(value) != 19 or value[10:] != " 00:00:00":
                    raise ValueError(f"{field} 格式错误：{value}")
    except json.JSONDecodeError as e:
        raise ValueError(f"JSON格式错误: {e}")

    return cleaned

# Utility function: 解析时间自然语言
# def parse_relative_date(relative: str) -> str:
#     """
#     将类似“昨天”、“今天”、“前天”解析为具体日期的字符串 (YYYY-MM-DD)
#     """
#     now = datetime.now()  # 获取当前时间
#     mapping = {
#         "今天": now,
#         "昨天": now - timedelta(days=1),
#         "前天": now - timedelta(days=2)
#     }
#     # 默认返回“未知”，只有识别到特定描述时才处理
#     return mapping.get(relative, "未知").strftime("%Y-%m-%d") if relative in mapping else "未知"


# 链式操作
extraction_chain = extraction_prompt | llm

def extract(classification_result, question):
    # 动态获取当前日期
    current_date = datetime.now().strftime("%Y-%m-%d")

    # 确保同时传递 `question` 和 `current_date`
    extraction_result = extraction_chain.invoke({"question": question, "current_date": current_date})

    # 清理模型输出
    cleaned_result = clean_model_output(extraction_result)
    extracted_data = ExtractedData.validate_output(cleaned_result)
    # try:
    #     extracted_data = ExtractedData.validate_output(cleaned_result)
    #     # 如果日期是相对时间（如“昨天”），进行解析
    #     if extracted_data["date"] in ["昨天", "今天", "前天"]:
    #         extracted_data["date"] = parse_relative_date(extracted_data["date"])
    # except ValueError as error:
    #     extracted_data = {"error": str(error)}

    return {
        "classification": classification_result,
        "extracted_data": extracted_data,
    }

# 测试代码
# classification_result = "轨迹查询"
# # example_question = "浙海盐货001昨天的航道轨迹"
# # example_question = "莲花号前天的航道轨迹"
# # example_question = "莲花号2025年1月20号的航道轨迹"
# example_question = "浙海盐货001从2025年1月20号到23号的航道轨迹"
# result = extract(classification_result, example_question)
# print(result)